import numpy as np
from zkl_aiutils_training import Resumable
from zkl_pyutils_fsspec import FsLike, resolve_child_fs

from .metrics import MetricOrValue, get_metric_scalar_value
from .metrics_logger import MetricsLogger
from .resumable_fs import FsResumeArgs


class ZklRecordsMetricsLogger(MetricsLogger, Resumable):
    def __init__(self, fs: FsLike | None = None):
        self._fs = fs

    @property
    def fs(self):
        return self._fs

    def on_resume(self, args: FsResumeArgs):
        if self._fs is None and args.fs is not None:
            self._fs = resolve_child_fs(args.fs, "records")

    def log(self, metrics: dict[str, MetricOrValue]):
        if self._fs is not None:
            for metric_name, metric in metrics.items():
                from zkl_aiutils_records import NumpySequenceRecordWriter
                metric_fs = resolve_child_fs(self._fs, metric_name)
                writer = NumpySequenceRecordWriter(metric_fs, shape=(), dtype=np.float32)
                writer.write(get_metric_scalar_value(metric))
